"""
数据预处理脚本 - 划分训练集、验证集、测试集
"""
import pandas as pd
from sklearn.model_selection import train_test_split
import os

def prepare_data(input_path, output_dir):
    """
    读取原始数据并划分为训练集、验证集、测试集

    参数:
        input_path: 原始CSV文件路径
        output_dir: 输出目录
    """
    print("📊 开始加载数据...")
    df = pd.read_csv(input_path)

    print(f"总数据量: {len(df)}")
    print(f"\n标签分布:")
    print(df['label'].value_counts())
    print(f"\nlabel 0 (安全): {(df['label']==0).sum()} ({(df['label']==0).sum()/len(df)*100:.2f}%)")
    print(f"label 1 (危险): {(df['label']==1).sum()} ({(df['label']==1).sum()/len(df)*100:.2f}%)")

    # 清洗数据：去除空值
    df = df.dropna()
    print(f"\n清洗后数据量: {len(df)}")

    # 划分数据集：70% 训练，15% 验证，15% 测试
    # 使用分层采样保持标签分布
    train_df, temp_df = train_test_split(
        df,
        test_size=0.3,
        random_state=42,
        stratify=df['label']
    )

    val_df, test_df = train_test_split(
        temp_df,
        test_size=0.5,
        random_state=42,
        stratify=temp_df['label']
    )

    print(f"\n✅ 数据集划分完成:")
    print(f"训练集: {len(train_df)} 条")
    print(f"验证集: {len(val_df)} 条")
    print(f"测试集: {len(test_df)} 条")

    # 保存数据集
    os.makedirs(output_dir, exist_ok=True)

    train_path = os.path.join(output_dir, 'train.csv')
    val_path = os.path.join(output_dir, 'val.csv')
    test_path = os.path.join(output_dir, 'test.csv')

    train_df.to_csv(train_path, index=False)
    val_df.to_csv(val_path, index=False)
    test_df.to_csv(test_path, index=False)

    print(f"\n💾 数据集已保存:")
    print(f"训练集: {train_path}")
    print(f"验证集: {val_path}")
    print(f"测试集: {test_path}")

    # 展示样本
    print(f"\n📝 训练集样本预览:")
    print(train_df.head(3))

if __name__ == "__main__":
    input_file = "../train/total_train.csv"
    output_directory = "./data"

    prepare_data(input_file, output_directory)
